package cn.sciento.core.algorithm.tree;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collector;
import java.util.stream.Collectors;

public class TreeBuilder {
    private TreeBuilder() throws IllegalAccessException {
        throw new IllegalAccessException();
    }

    public static <P, T extends Child<T>> List<T> buildTree(List<T> objList, Node<P, T> nodeOperation) {
        Map<P, T> objMap = (Map<P, T>) objList.stream().collect(Collectors.toMap(nodeOperation::getKey, Function.identity()));
        Map<P, List<T>> objGroupMap = new HashMap<>(16);
        objMap.forEach((key, value) -> {
            P parent;
            if (objMap.containsKey(nodeOperation.getParentKey(value))) {
                parent = nodeOperation.getParentKey(value);
            } else {
                parent = null;
            }
            if (objGroupMap.containsKey(parent)) {
                ((List) objGroupMap.get(parent)).add(value);
            } else {
                List<T> list = new ArrayList<>();
                list.add((T) value);
                objGroupMap.put(parent, list);
            }
        });
        return recursiveBuildTree(objGroupMap, null, new ArrayList<>(), nodeOperation);
    }

    public static <P, T extends Child<T>> List<T> buildTree(List<T> objList, P rootKey, Node<P, T> nodeOperation) {
        return buildTree(objList, rootKey, new ArrayList<>(), nodeOperation);
    }

    public static <P, T extends Child<T>> List<T> buildTree(List<T> objList, P rootKey, List<T> list, Node<P, T> nodeOperation) {
        return recursiveBuildTree((Map<P, List<T>>) objList.stream().collect(groupingByWithNullKeys(nodeOperation::getParentKey)), rootKey, list, nodeOperation);
    }

    public static <P, T extends Child<T>> List<T> buildTree(List<T> objList, P rootKey, Key<P, T> key, ParentKey<P, T> parentKey) {
        return buildTree(objList, rootKey, new ArrayList<>(), key, parentKey);
    }

    public static <P, T extends Child<T>> List<T> buildTree(List<T> objList, P rootKey, List<T> list, Key<P, T> key, ParentKey<P, T> parentKey) {
        return recursiveBuildTree((Map<P, List<T>>) objList.stream().collect(groupingByWithNullKeys(parentKey::getParentKey)), rootKey, list, key, parentKey);
    }

    public static <P, T extends Child<T>> List<T> recursiveBuildTree(Map<P, List<T>> map, P parentId, List<T> list, Node<P, T> nodeOperation) {
        if (map.containsKey(parentId)) {
            List<T> parentList = map.get(parentId);
            for (Child child : parentList) {
                if (map.containsKey(nodeOperation.getKey((T) child)))
                    child.addChildren(recursiveBuildTree(map, nodeOperation.getKey((T) child), new ArrayList<>(), nodeOperation));
                list.add((T) child);
            }
        }
        return list;
    }

    public static <P, T extends Child<T>> List<T> recursiveBuildTree(Map<P, List<T>> map, P parentId, List<T> list, Key<P, T> key, ParentKey<P, T> parentKey) {
        if (map.containsKey(parentId)) {
            List<T> parentList = map.get(parentId);
            for (Child child : parentList) {
                if (map.containsKey(key.getKey((T) child)))
                    child.addChildren(recursiveBuildTree(map, key.getKey((T) child), new ArrayList<>(), key, parentKey));
                list.add((T) child);
            }
        }
        return list;
    }

    public static <T, A> Collector<T, ?, Map<A, List<T>>> groupingByWithNullKeys(Function<? super T, ? extends A> classifier) {
        return Collectors.toMap(classifier, Collections::singletonList, (oldList, newEl) -> {
            List<T> newList = new ArrayList<>(oldList.size() + 1);
            newList.addAll(oldList);
            newList.addAll(newEl);
            return newList;
        });
    }
}
